# main_window.py
import sys
import os
import qdarkstyle
from dotenv import load_dotenv
from openai import OpenAI
from PySide6.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QTextEdit, 
    QPushButton, QLabel, QComboBox, QCheckBox, QHBoxLayout, QFrame
)
from PySide6.QtCore import QThread, Signal, Qt, QPoint
from PySide6.QtGui import QMouseEvent, QFont, QFontDatabase

# [移除] 不再需要导入不存在的库

# --- 后台API请求线程 ---
class ApiWorker(QThread):
    """
    在一个独立的线程中运行API请求，以防止UI冻结。
    支持流式响应，实时将数据块通过信号发送回主线程。
    """
    data_received = Signal(str)  # 定义一个信号，用于发送接收到的数据
    finished = Signal()          # 定义一个信号，在任务完成时发送

    def __init__(self, client, model, messages):
        super().__init__()
        self.client = client
        self.model = model
        self.messages = messages
        self._is_running = True

    def run(self):
        """线程执行的入口点"""
        try:
            stream = self.client.chat.completions.create(
                model=self.model,
                messages=self.messages,
                stream=True,
                temperature=0.7,
            )
            for chunk in stream:
                if not self._is_running:
                    break
                content = chunk.choices[0].delta.content
                if content:
                    self.data_received.emit(content)
        except Exception as e:
            self.data_received.emit(f"\nAPI请求出错: {e}")
        finally:
            self.finished.emit()

    def stop(self):
        self._is_running = False

# --- 主窗口 ---
class MainWindow(QWidget):
    """
    应用的主窗口，一个可拖动的无边框悬浮窗。
    """
    def __init__(self):
        super().__init__()
        self.api_worker = None
        self._drag_pos = QPoint()

        self.load_env()
        self.setup_font()
        self.init_ui()
        self.apply_stylesheet()
        
    def load_env(self):
        """加载环境变量，但不再直接初始化客户端"""
        load_dotenv()
        self.deepseek_api_key = os.getenv("DEEPSEEK_API_KEY")

    def setup_font(self):
        """
        [修复] 优先从本地文件加载字体，如果失败则回退到系统字体。
        这是解决乱码的最可靠方法。
        """
        font_path = os.path.join("assets", "fonts", "NotoSansSC-Regular.ttf")
        if os.path.exists(font_path):
            font_id = QFontDatabase.addApplicationFont(font_path)
            if font_id != -1:
                family = QFontDatabase.applicationFontFamilies(font_id)[0]
                QApplication.instance().setFont(QFont(family, 10))
                print(f"成功从文件加载并设置全局字体: {family}")
                return

        print(f"警告：未在 '{font_path}' 找到字体文件。")
        print("将尝试在系统中搜索备用字体...")
        
        font_families = [
            "Microsoft YaHei", "PingFang SC", "Noto Sans CJK SC", 
            "WenQuanYi Micro Hei", "SimSun",
        ]
        db = QFontDatabase()
        for family in font_families:
            if family in db.families():
                QApplication.instance().setFont(QFont(family, 10))
                print(f"成功找到并设置备用系统字体: {family}")
                return
        
        print("警告：未找到任何推荐的中文字体，界面文本可能出现乱码。")

    def init_ui(self):
        """初始化用户界面"""
        self.setWindowFlags(
            Qt.WindowType.FramelessWindowHint |
            Qt.WindowType.Tool |
            Qt.WindowType.WindowStaysOnTopHint
        )
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)
        self.resize(400, 600)

        # 主布局
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(15, 15, 15, 15)

        # 创建一个带边框的Frame作为背景
        bg_frame = QFrame(self)
        bg_frame.setObjectName("backgroundFrame")
        main_layout.addWidget(bg_frame)
        
        layout = QVBoxLayout(bg_frame)
        layout.setSpacing(15)

        # [修改] 移除API选择框，直接显示当前模型
        model_label = QLabel("模型: DeepSeek (deepseek-v2)")
        model_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(model_label)

        # 输入框
        layout.addWidget(QLabel("原始Prompt:"))
        self.input_text = QTextEdit()
        self.input_text.setPlaceholderText("在这里输入您想要优化的Prompt...")
        self.input_text.setMinimumHeight(150)
        layout.addWidget(self.input_text)

        # 优化按钮
        self.optimize_button = QPushButton("开始优化")
        self.optimize_button.clicked.connect(self.start_optimization)
        layout.addWidget(self.optimize_button)

        # 输出框
        output_header_layout = QHBoxLayout()
        output_header_layout.addWidget(QLabel("优化结果:"))
        output_header_layout.addStretch()
        self.copy_button = QPushButton("复制")
        self.copy_button.clicked.connect(self.copy_output_text)
        output_header_layout.addWidget(self.copy_button)
        layout.addLayout(output_header_layout)

        self.output_text = QTextEdit()
        self.output_text.setAcceptDrops(False)  # 禁用拖放功能
        self.output_text.setMinimumHeight(200)
        layout.addWidget(self.output_text)

        # --- 控制面板 ---
        panel_layout = QHBoxLayout()
        
        self.stay_on_top_checkbox = QCheckBox("保持在最前")
        self.stay_on_top_checkbox.setChecked(True)
        self.stay_on_top_checkbox.stateChanged.connect(self.toggle_stay_on_top)
        panel_layout.addWidget(self.stay_on_top_checkbox)

        panel_layout.addStretch() # 添加伸缩项

        self.quit_button = QPushButton("退出")
        self.quit_button.clicked.connect(QApplication.instance().quit)
        panel_layout.addWidget(self.quit_button)
        
        layout.addLayout(panel_layout)

    def apply_stylesheet(self):
        """应用QDarkStyle样式并自定义"""
        self.setStyleSheet(qdarkstyle.load_stylesheet(qt_api='pyside6'))
        custom_style = """
            #backgroundFrame {
                background-color: #2d3436;
                border-radius: 10px;
                border: 1px solid #7D7D7D;
            }
            QTextEdit, QComboBox {
                border-radius: 5px;
            }
            QPushButton {
                padding: 8px;
                border-radius: 5px;
            }
        """
        self.setStyleSheet(self.styleSheet() + custom_style)

    def start_optimization(self):
        """开始优化过程，按需创建API客户端"""
        if self.api_worker and self.api_worker.isRunning():
             self.api_worker.stop()
             self.optimize_button.setText("开始优化")
             return

        self.output_text.clear()
        
        prompt = self.input_text.toPlainText()
        
        if not prompt.strip():
            self.output_text.setText("错误：请输入需要优化的Prompt。")
            return

        self.optimize_button.setText("停止...")
        
        system_prompt = "你是一位Prompt工程大师。你的任务是接收用户输入的Prompt，然后将其改写得更清晰、更具体、更高效，以便大型语言模型能更好地理解和执行。请直接返回优化后的Prompt，不要包含任何额外的解释或说明文字。"
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ]

        # [修改] 硬编码为使用DeepSeek
        if not self.deepseek_api_key:
            self.output_text.setText("错误：未检测到 DEEPSEEK_API_KEY。\n请在 .env 文件中配置您的API密钥。")
            self.optimize_button.setText("开始优化")
            return
        client = OpenAI(api_key=self.deepseek_api_key, base_url="https://api.deepseek.com/v1")
        model = "deepseek-chat"

        self.api_worker = ApiWorker(client, model, messages)
        self.api_worker.data_received.connect(self.update_output)
        self.api_worker.finished.connect(lambda: self.optimize_button.setText("开始优化"))
        self.api_worker.start()

    def update_output(self, text):
        """将从线程接收到的数据显示在输出框中"""
        self.output_text.insertPlainText(text)

    def toggle_stay_on_top(self, state):
        """
        [修复] 回退到最可靠的X11兼容置顶方案。
        对于Wayland环境，需要使用外部工具。
        """
        self.hide()
        # 包含了之前尝试过最可能生效的Tool标志
        base_flags = Qt.WindowType.FramelessWindowHint | Qt.WindowType.Tool
        if state == Qt.CheckState.Checked.value:
            self.setWindowFlags(base_flags | Qt.WindowType.WindowStaysOnTopHint)
        else:
            self.setWindowFlags(base_flags)
        self.show()
        self.raise_()
        self.activateWindow()

    def copy_output_text(self):
        """将优化后的文本复制到剪贴板。"""
        QApplication.clipboard().setText(self.output_text.toPlainText())

    # --- 窗口拖动逻辑 ---
    def mousePressEvent(self, event: QMouseEvent):
        if event.button() == Qt.MouseButton.LeftButton:
            self._drag_pos = event.globalPosition().toPoint() - self.frameGeometry().topLeft()
            event.accept()

    def mouseMoveEvent(self, event: QMouseEvent):
        if event.buttons() == Qt.MouseButton.LeftButton:
            self.move(event.globalPosition().toPoint() - self._drag_pos)
            event.accept()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())
